🌊 WaveSpeed AI: High-Performance Multimodal Generation

#!pip install wavespeed python-dotenv requests -q

import os
import wavespeed
from dotenv import load_dotenv
from wavespeed import Client

# Load environment variables from .env file
load_dotenv()

api_key = os.getenv("WAVESPEED_API_KEY")
client = Client(api_key=api_key)

if not api_key:
    raise ValueError("API Key not found! Please check your .env file.")

print("✅ API Key loaded successfully.")

from IPython.display import Image, display
import requests

def download_and_display(url, filename="output.png"):
    """Helper to download and display the generated media."""
    response = requests.get(url)
    if response.status_code == 200:
        with open(filename, 'wb') as f:
            f.write(response.content)
        display(Image(filename=filename, width=600))
        print(f"Saved to {filename}")
    else:
        print(f"Failed to download media: {response.status_code}")

# ------------------------------
# Generate an Image
# ------------------------------

print("🎨 Generating image with Flux...")

try:
    # Note: Model IDs may change. Check wavespeed.ai/models for the latest list.
    # Common IDs: 'wavespeed-ai/flux-dev', 'black-forest-labs/flux-1-dev'
    result = client.run(
        "wavespeed-ai/flux-dev", 
        {
            "prompt": "A cyberpunk street food vendor in Tokyo, neon lights, rain reflections, highly detailed, 8k",
            "aspect_ratio": "16:9",
            "guidance_scale": 3.5, # This controls how strictly the AI adheres to your text prompt; wiht lower value, the AI takes more artistic liberties.
            "num_inference_steps": 20 # This is the number of "passes" the AI makes to clear up the image from noise to a final picture; with lower value, faster generation, but the image might look blurry
        }
    )

    # The output is usually a list of URLs
    image_url = result["outputs"][0]["url"] if isinstance(result["outputs"][0], dict) else result["outputs"][0]
    print(f"Generated URL: {image_url}")
    
    download_and_display(image_url, "flux_result.png")

except Exception as e:
    print(f"Error generating image: {e}")

from IPython.display import Video

# We will use the image URL generated in the previous step.
# If you skipped that step, replace this with any public image URL.
input_image_url = image_url 

print("🎥 Generating video from image... (This may take a minute)")

try:
    # Using Wan 2.1 (Example ID - verify exact ID on the platform)
    # Other popular options: 'kwaivgi/kling-v2.6-i2v'
    video_result = client.run(
        "wan-ai/wan-2.1-i2v-720p", 
        {
            "image_url": input_image_url,
            "prompt": "The vendor cooks food, steam rising, neon lights flickering, cinematic slow motion",
            "negative_prompt": "distortion, blurry, shaking camera",
            "duration": 5  # Duration in seconds (depends on model capabilities)
        }
    )

    video_output_url = video_result["outputs"][0]["url"] if isinstance(video_result["outputs"][0], dict) else video_result["outputs"][0]
    print(f"Video URL: {video_output_url}")

    # Display Video
    display(Video(url=video_output_url, width=600))

except Exception as e:
    print(f"Error generating video: {e}")

print("🎨 Generating stylized image with LoRA...")

try:
    lora_result = wavespeed.run(
        "wavespeed-ai/flux-dev",
        {
            "prompt": "A cute robot holding a flower, papercut style, intricate details",
            # LoRA configuration syntax depends on the specific model support
            # WaveSpeed often allows passing a 'loras' list or embedding it in the prompt for some models.
            "loras": [
                {
                    "path": "TheMisto/Misto-Flux-Papercut", # Example HF path
                    "scale": 0.8  # Strength of the LoRA
                }
            ],
            "aspect_ratio": "1:1",
            "guidance_scale": 4.0
        }
    )

    lora_url = lora_result["outputs"][0]["url"] if isinstance(lora_result["outputs"][0], dict) else lora_result["outputs"][0]
    download_and_display(lora_url, "lora_result.png")

except Exception as e:
    print(f"Error with LoRA generation: {e}")

import concurrent.futures

prompts = [
    "A futuristic city on Mars",
    "An underwater civilization"
]

def generate_batch(prompt):
    print(f"Starting: {prompt}...")
    res = wavespeed.run("wavespeed-ai/flux-dev", {"prompt": prompt})
    return res["outputs"][0]

# Run in parallel threads
# Note: Ensure your API tier supports concurrent requests
try:
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        results = list(executor.map(generate_batch, prompts))
    
    print("Batch generation complete!")
    for i, url in enumerate(results):
        print(f"Result {i+1}: {url}")
except Exception as e:
    print(f"Batch processing error: {e}")